﻿unit RateLimitMiddleware;

interface

uses
  System.SysUtils, System.Generics.Collections, System.DateUtils, System.SyncObjs, System.Classes,
  Sparkle.HttpServer.Module, Sparkle.HttpServer.Context;

type
  TRateLimitInfo = record
    RequestCount: Integer;
    WindowStart: TDateTime;
  end;

  TRateLimitMiddleware = class(THttpServerMiddleware)
  private
    FMaxRequests: Integer;
    FWindowMinutes: Integer;
    FClientLimits: TDictionary<string, TRateLimitInfo>;
    FLock: TCriticalSection;

    function GetClientKey(Context: THttpServerContext): string;
    function IsRateLimited(const ClientKey: string): Boolean;
    procedure UpdateRateLimit(const ClientKey: string);
  protected
    procedure ProcessRequest(Context: THttpServerContext; Next: THttpServerProc); override;
  public
    constructor Create(MaxRequests: Integer = 100; WindowMinutes: Integer = 1);
    destructor Destroy; override;
  end;

implementation

uses
  System.NetEncoding;

{ TRateLimitMiddleware }

constructor TRateLimitMiddleware.Create(MaxRequests, WindowMinutes: Integer);
begin
  inherited Create;
  FMaxRequests := MaxRequests;
  FWindowMinutes := WindowMinutes;
  FClientLimits := TDictionary<string, TRateLimitInfo>.Create;
  FLock := TCriticalSection.Create;
end;

destructor TRateLimitMiddleware.Destroy;
begin
  FLock.Free;
  FClientLimits.Free;
  inherited;
end;

function TRateLimitMiddleware.GetClientKey(Context: THttpServerContext): string;
begin
  Result := Context.Request.RemoteIp;
end;

function TRateLimitMiddleware.IsRateLimited(const ClientKey: string): Boolean;
var
  LimitInfo: TRateLimitInfo;
  CurrentWindow: TDateTime;
begin
  FLock.Enter;
  try
    CurrentWindow := Trunc(Now * (24 * 60 / FWindowMinutes)) / (24 * 60 / FWindowMinutes);

    if FClientLimits.TryGetValue(ClientKey, LimitInfo) then
    begin
      if LimitInfo.WindowStart < CurrentWindow then
      begin
        // 新时间窗口，重置计数
        LimitInfo.RequestCount := 1;
        LimitInfo.WindowStart := CurrentWindow;
        FClientLimits[ClientKey] := LimitInfo;
        Result := False;
      end
      else
      begin
        // 检查是否超过限制
        Result := LimitInfo.RequestCount >= FMaxRequests;
        if not Result then
        begin
          Inc(LimitInfo.RequestCount);
          FClientLimits[ClientKey] := LimitInfo;
        end;
      end;
    end
    else
    begin
      // 新客户端
      LimitInfo.RequestCount := 1;
      LimitInfo.WindowStart := CurrentWindow;
      FClientLimits.Add(ClientKey, LimitInfo);
      Result := False;
    end;
  finally
    FLock.Leave;
  end;
end;

procedure TRateLimitMiddleware.UpdateRateLimit(const ClientKey: string);
begin
  // 在IsRateLimited中已经更新，这里不需要额外操作
end;

procedure TRateLimitMiddleware.ProcessRequest(Context: THttpServerContext;
  Next: THttpServerProc);
var
  ClientKey: string;
begin
  try
    ClientKey := GetClientKey(Context);

    if IsRateLimited(ClientKey) then
    begin
      // 设置限流响应头
      Context.Response.StatusCode := 429; // Too Many Requests
      Context.Response.Headers.SetValue('Retry-After', IntToStr(FWindowMinutes * 60));
      Context.Response.ContentType := 'application/json; charset=utf-8';

      // 使用 Close 方法发送响应
      Context.Response.Close(TEncoding.UTF8.GetBytes(
        '{"success": false, "code": 429, "message": "Rate limit exceeded. Too many requests, please try again later."}'));

      Exit;
    end;

    // 如果没有被限流，继续处理
    Next(Context);

  except
    on E: Exception do
    begin
      // 发生异常时，继续处理请求而不是中断
      try
        Next(Context);
      except
        // 忽略第二次异常
      end;
    end;
  end;
end;

end.
